Google ColabからHugging FaceのモデルをVertex AIへデプロイしてみた

皆さん、こんにちは。

クルトンです。

今回のエントリーは『機械学習チームアドベントカレンダー2022』の20日目の記事になります。

前日(12/19)の記事は『Amazon SageMaker Data Wranglerに再入門する（2022年12月版）その１』でした。 入門記事として料金や機能の概要、実際の操作について書かれていますので、ぜひご覧ください！

今回は、こちらの記事を参考に、Hugging FaceモデルをVertex AIにデプロイしてみます。

実行環境はGoogle ColabのCPUです。また、Google Cloudでのアカウントが既に作成完了しており、プロジェクトが作られている事を前提としています。

環境セットアップ

まずはGoogle Cloudのリソースへプログラムからアクセスできるようにしていきます。

from google.colab import auth
auth.authenticate_user()

このコードを実行すると、次のような画面が出てきます。

画像の「許可」のボタンをクリックしてから、出てきた画面でGoogle Cloudのアカウントを選択。 さらにもう一度「許可」のボタンをクリックするとColabからGoogle Cloudへのリソースにアクセスする前準備は完了です。

バケット

Vertex AIにモデルをデプロイする過程で、Cloud Storageにモデルをアップロードする必要があります。

既にバケットを用意している場合は、そのバケット名を記入ください。作成がまだの場合はgsutilコマンドを叩いて作成しますので、作りたいバケット名を記入してください。

今回はasia-northeast1で全て実行していますが、今回使うサービスが動くリージョンであれば他のリージョンに書き換えても動くかと思います。

GCS_BUCKET = "gs://<ここをお好きな名前に変更してください>"
REGION = "asia-northeast1"

バケットがない場合にコマンドから作成

バケットがない場合はgsutilコマンドを実行して作成します。その前の段階としてコマンドで作成できるようにアクセスの許可を与えていきます。

まずはログインをします。

!gcloud auth login

Do you want to continue (Y/n)?と実行結果の部分で表示されますので、こちらに"Y"を入力します。

そうするとURLが出てきますので、そちらをクリックします。

クリック後に再び、アカウントを選択する画面が出てきますので、そちらで先ほど選択したものと同じアカウントを選択します。

アカウント選択後に次のような画面が出てきますので、画面下部のコピーを選択します。

Google Colabの画面に戻り、コピーしたテキストをEnter authorization codeへ貼り付けてEnterキーを押してください。これで、Google Colabで該当のアカウントに対してコマンドを叩いての操作が出来るようになりました。

次にprojectの設定をしておきます。

!gcloud config set project <自分のプロジェクト名に書き換え>

gcloudコマンドの内容が反映されているか気になる方は次のコマンドでご確認ください。

!gcloud config list

次のコマンドでバケットを作成します。

!gsutil mb -l $REGION $GCS_BUCKET

Google Cloudコンソール上でも作成できている事が確認できますので、気になる方はそちらご確認ください。

必要なモジュールを調整

Google Cloudへのモデルデプロイなどに必要なモジュールをアップデートとインストールをしておきます。

!pip install --upgrade google-cloud-aiplatform transformers -q

次のモジュールをuninstallしておきます。

!pip uninstall shapely -y

こちらのモジュールを削除しておく理由ですが、後ほどimportするgapicでImportError : cannot import name 'WKBWriter' from 'shapely.geos'が発生するためです。再度入れ直しをした場合でもエラーが直らなかったので削除しておきます。

ランタイムの再起動

ここでランタイムの再起動をする必要があります。そうしないと後ほどimportするgapicモジュールでContextualVersionConflictというエラーが出てきます。これは、google-cloud-bigqueryのバージョンによるエラーです。3.3.6のバージョンが元からColabには入っているのですが、gapicで使用するバージョンは2.34.4と、想定されているバージョンよりも高いために起こるものです。(より正確にはバージョンが3以上のものはエラーが起きるようです。)

ランタイム再起動はColabの画面上にあるランタイムのタブから「ランタイムを再起動」から実行できます。

必要なモジュールを調整

ここで、以降で必要となるものを一括でimportします。

from transformers import TFViTForImageClassification
import tensorflow as tf
import os
from google.cloud.aiplatform import gapic as aip

バージョンについて、tensorflowは2.9.2、transformersは4.25.1です。

たまにここでtransformersなどのimportでエラーが出る場合があります。その時はpip installとpip uninstallやランタイムの再起動をやり直す事でエラーが解決する事がありますので、もう一度実行してみてください。

モデル作成

次のコードでHugging Faceから学習済みモデルをロードします。

LOCAL_MODEL_DIR = "vit"
model = TFViTForImageClassification.from_pretrained("google/vit-base-patch16-224")

# Colabの一時ストレージに保存したい場合は実行
model.save_pretrained(LOCAL_MODEL_DIR, saved_model=True) 

次のコマンドを使うとSaved型のモデルの詳細について知る事ができます。今回ロードしたモデルを確認できます。

!saved_model_cli show --dir {LOCAL_MODEL_DIR}/saved_model/1 --all

Cloud Storageにモデルをアップロードします。

tf.saved_model.save(
    model,
    os.path.join(GCS_BUCKET, LOCAL_MODEL_DIR),
)

バケット名など事前定義した場所へ、変数modelへロードしたモデルをエクスポートします。

Vertex AIのエンドポイントを作成してモデルをデプロイ

ここでエンドポイント作成やモデルデプロイに必要なClientのインスタンスを作成します。

PROJECT_ID = <自身のプロジェクトIDを入力ください>
API_ENDPOINT = f"{REGION}-aiplatform.googleapis.com"
PARENT = f"projects/{PROJECT_ID}/locations/{REGION}"

client_options = {"api_endpoint": API_ENDPOINT}
model_service_client = aip.ModelServiceClient(client_options=client_options)
endpoint_service_client = aip.EndpointServiceClient(client_options=client_options)

Vertex AI Model Registryへアップロード

次のコードでモデルをアップロードします。 変数tf28_gpu_model_dictにあるcontainer_specはArtifact Registryというコンテナを管理するGoogle Cloudのサービスに上げているコンテナを指定します。

tf28_gpu_model_dict = {
    "display_name": "ViT Base TF2.8 GPU model",
    "artifact_uri": f"{GCS_BUCKET}/{LOCAL_MODEL_DIR}",
    "container_spec": {
        "image_uri": "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-8:latest",
    },
}

tf28_gpu_model = (
    model_service_client.upload_model(parent=PARENT, model=tf28_gpu_model_dict)
    .result(timeout=180)
    .model
)
tf28_gpu_model

アップロードが完了するとVertex AI Model Registryでモデルを確認出来ます。アップロードに使用したコードの出力結果と同じモデルIDである事を確認出来ます。

エンドポイントの作成

次のコードでエンドポイントの作成をします。

tf28_gpu_endpoint_dict = {
    "display_name": "ViT Base TF2.8 GPU endpoint",
}
tf28_gpu_endpoint = (
    endpoint_service_client.create_endpoint(
        parent=PARENT, endpoint=tf28_gpu_endpoint_dict
    )
    .result(timeout=300)
    .name
)
tf28_gpu_endpoint

エンドポイントへのモデルをデプロイ

次のコードでデプロイしたモデルが使用するCompute Engineのリソースを設定します。

DEPLOY_COMPUTE = "n1-standard-8"
DEPLOY_GPU = aip.AcceleratorType.NVIDIA_TESLA_T4

次にエンドポイントへモデルのアップロードを行います。

tf28_gpu_deployed_model_dict = {
    "model": tf28_gpu_model,
    "display_name": "ViT Base TF2.8 GPU deployed model",
    "dedicated_resources": {
        "min_replica_count": 1,
        "max_replica_count": 1,
        "machine_spec": {
            "machine_type": DEPLOY_COMPUTE,
            "accelerator_type": DEPLOY_GPU,
            "accelerator_count": 1,
        },
    },
}

tf28_gpu_deployed_model = endpoint_service_client.deploy_model(
    endpoint=tf28_gpu_endpoint,
    deployed_model=tf28_gpu_deployed_model_dict,
    traffic_split={"0": 100},
).result()
tf28_gpu_deployed_model

エンドポイントへのモデルのデプロイは、Vertex AI Model Registryへのアップロードと比べると時間が掛かります。 自分の実行時には、6~9分ほど掛かるケースが多かったです。

クリーンアップ

ここまでで作ってきたリソースを、今後使用しない場合は削除しましよう。 次のコードを実行します。

def cleanup(endpoint, model_name, deployed_model_id):
    response = endpoint_service_client.undeploy_model(
        endpoint=endpoint, deployed_model_id=deployed_model_id
    )
    print("running undeploy_model operation:", response.operation.name)
    print(response.result())

    response = endpoint_service_client.delete_endpoint(name=endpoint)
    print("running delete_endpoint operation:", response.operation.name)
    print(response.result())

    response = model_service_client.delete_model(name=model_name)
    print("running delete_model operation:", response.operation.name)
    print(response.result())


cleanup(tf28_gpu_endpoint, tf28_gpu_model, tf28_gpu_deployed_model.deployed_model.id)

これでVertex AIにアップロードしたモデル、エンドポイント、エンドポイントにデプロイしたモデルを削除できます。

バケットに上げたモデルを削除します。

!gsutil rm -r $GCS_BUCKET

終わりに

今回はHugging Faceモデルをデプロイしてみました！ Artifact Registryへカスタムコンテナを上げると同じようにColabからデプロイなどが可能になります。

エンドポイントへデプロイ時にtrafficの調整も可能なので便利そうですね！

今回はここまで。

それでは、また！

参考にしたサイト